# environment.py
import gym
from gym import spaces
import numpy as np
import pygame
import math
import time
import random
import cv2
from resourceloader import ResourcesLoader
from collections import deque


class rm_simple_simulator(gym.Env):
    def __init__(self, render_mode="human"):
        super(rm_simple_simulator, self).__init__()
        self.loader = ResourcesLoader()
        # self.loader.params_printer()

        # Robot types and teams
        self.robot_types = ["infantry", "hero", "sentry"]
        self.teams = ["red", "blue"]
        self.num_robots = 3 * 2  # 3v3

        # Action space: per robot [velocity (0-1), rotation (-1 to 1), attack (0/1)]
        self.action_space = spaces.Box(
            low=np.array([0, -1, 0] * self.num_robots, dtype=np.float32),
            high=np.array([1, 1, 1] * self.num_robots, dtype=np.float32),
            dtype=np.float32,
        )
        # Observation: RGB vision image per robot? But for simplicity, global state or single vision? Doc says per robot vision.
        # For gym, observation is dict or array; here, we'll make it a stack of vision images, but to simplify, single global RGB for now.
        # Doc: observation is RGB image for robot vision.
        # But for multi-agent, need per agent. For simplicity, assume single agent view, but since 3v3, perhaps global.
        # To match, let's assume observation is a tuple of visions.
        vision_size = int(
            self.loader.robot_vision * 2 * self.loader.resolution
        )  # diameter
        self.output_mat_size = 128
        self.observation_space = spaces.Dict(
            {
                "visions": spaces.Box(
                    low=0,
                    high=255,
                    shape=(self.num_robots, vision_size, vision_size, 3),
                    dtype=np.uint8,
                ),
                "infos": spaces.Dict(
                    {
                        "time": spaces.Box(
                            low=0,
                            high=float(self.loader.time),
                            shape=(),
                            dtype=np.float32,
                        ),
                        "hp": spaces.Box(
                            low=0,
                            high=1000.0,
                            shape=(self.num_robots,),
                            dtype=np.float32,
                        ),
                        "bullet": spaces.Box(
                            low=0, high=800.0, shape=(self.num_robots,), dtype=np.int32
                        ),
                        "attack_cooldown": spaces.Box(
                            low=0, high=10.0, shape=(self.num_robots,), dtype=np.float32
                        ),
                        "reborn_timer": spaces.Box(
                            low=0,
                            high=300.0,
                            shape=(self.num_robots,),
                            dtype=np.float32,
                        ),
                        "scores": spaces.Dict(
                            {
                                "red": spaces.Box(
                                    low=0, high=1000.0, shape=(), dtype=np.float32
                                ),
                                "blue": spaces.Box(
                                    low=0, high=1000.0, shape=(), dtype=np.float32
                                ),
                            }
                        ),
                    }
                ),
            }
        )

        # State
        self.map_vis = self.loader.map_vis
        self.map_walkable = np.logical_not(self.loader.map_obstacles)
        self.map_obstacles = self.loader.map_obstacles
        self.map_bonus = self.loader.map_bonus
        self.map_red_ammo = self.loader.map_red_ammo_zone
        self.map_blue_ammo = self.loader.map_blue_ammo_zone
        self.h, self.w = self.map_vis.shape[:2]

        # Robot params
        self.robot_size = self.loader.robot_real_width  # meters side length from yaml
        self.robot_radius_px = int(self.robot_size * self.loader.resolution / 2)
        self.max_velocity = 0.8  # m/s assume
        self.max_rotation = np.pi / 2  # rad/s assume
        self.vision_size = int(2 * self.loader.robot_vision * self.loader.resolution)

        # Scores
        self.scores = {"red": 0, "blue": 0}
        self.current_time = 0
        self.dt = 1 / 10

        # Additional reward parameters for denser rewards
        self.damage_reward = 1.0  # Base reward per fraction of max HP damaged
        self.supply_reward = 0.0  # Disabled flat reward, use gain-based instead
        self.survival_reward = 0.01  # Small reward per dt for staying alive
        self.move_reward = 0.01  # Small reward per dt for moving

        # Use original YAML values for sparse rewards to guide RL properly
        # reward_kill: 20, reward_center_occupy: 1, victory_reward: 200
        # These provide balance: kills dominant (20/episode event), occupy steady (1/s ~0.1/step), victory terminal (200)

        # Robots state: list of dicts {team, type, pos (x,y), rot, hp, reborn_timer, attack_cooldown, alive}
        self.robots = self._init_robots()

        # Scene image for visions
        self.scene_image = None

        # Visualization states for effects
        self.firing_robots = []  # list of indices that fired this frame
        self.damaged_robots = []  # list of indices that took damage this frame

        # Render
        self.render_mode = render_mode
        self.screen = None
        self.clock = None
        self.font = None  # Will be initialized in render

        # Gym metadata
        self.metadata = {"render_modes": ["human"], "render_fps": 60}

    def _init_robots(self):
        robots = []
        starts = {
            "red": {
                "infantry": (
                    self.loader.red_infantry_start[0]
                    if len(self.loader.red_infantry_start) > 0
                    else np.array([100, 100])
                ),
                "hero": (
                    self.loader.red_hero_start[0]
                    if len(self.loader.red_hero_start) > 0
                    else np.array([100, 100])
                ),
                "sentry": (
                    self.loader.red_sentry_start[0]
                    if len(self.loader.red_sentry_start) > 0
                    else np.array([100, 100])
                ),
            },
            "blue": {
                "infantry": (
                    self.loader.blue_infantry_start[0]
                    if len(self.loader.blue_infantry_start) > 0
                    else np.array([self.w - 100, self.h - 100])
                ),
                "hero": (
                    self.loader.blue_hero_start[0]
                    if len(self.loader.blue_hero_start) > 0
                    else np.array([self.w - 100, self.h - 100])
                ),
                "sentry": (
                    self.loader.blue_sentry_start[0]
                    if len(self.loader.blue_sentry_start) > 0
                    else np.array([self.w - 100, self.h - 100])
                ),
            },
        }
        for team in self.teams:
            for rtype in self.robot_types:
                hp = (
                    self.loader.infantry_hp
                    if rtype == "infantry"
                    else (
                        self.loader.hero_hp
                        if rtype == "hero"
                        else self.loader.sentry_hp
                    )
                )
                reborn = (
                    self.loader.infantry_reborn_time
                    if rtype == "infantry"
                    else (
                        self.loader.hero_reborn_time
                        if rtype == "hero"
                        else self.loader.sentry_reborn_time
                    )
                )
                attack_freq = (
                    self.loader.infantry_attack_frequency
                    if rtype == "infantry"
                    else (
                        self.loader.hero_attack_frequency
                        if rtype == "hero"
                        else self.loader.sentry_attack_frequency
                    )
                )
                attack_dmg = (
                    self.loader.infantry_attack
                    if rtype == "infantry"
                    else (
                        self.loader.hero_attack
                        if rtype == "hero"
                        else self.loader.sentry_attack
                    )
                )
                bullet_max = (
                    self.loader.infantry_bullet_max
                    if rtype == "infantry"
                    else (
                        self.loader.hero_bullet_max
                        if rtype == "hero"
                        else self.loader.sentry_bullet_max
                    )
                )
                pos = starts[team][rtype].astype(float)
                # Clamp position to map bounds
                pos[0] = np.clip(pos[0], 0, self.w - 1)
                pos[1] = np.clip(pos[1], 0, self.h - 1)
                robots.append(
                    {
                        "team": team,
                        "type": rtype,
                        "pos": pos,
                        "rot": 0,
                        "hp": hp,
                        "max_hp": hp,
                        "reborn_timer": 0,
                        "attack_cooldown": 0,
                        "attack_freq": attack_freq,
                        "attack_dmg": attack_dmg,
                        "bullet": bullet_max,
                        "max_bullet": bullet_max,
                        "alive": True,
                        "reborn_time": reborn,  # Store reborn_time for each robot
                    }
                )
        random.shuffle(robots)  # for action indexing
        return robots

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.scores = {"red": 0, "blue": 0}
        self.current_time = 0
        self.robots = self._init_robots()
        self.firing_robots = []
        self.damaged_robots = []
        self._draw_scene()
        obs = self._get_obs()
        return obs, {}

    def step(self, actions):
        # Clear previous frame effects
        self.firing_robots = []
        self.damaged_robots = []

        # Actions: flattened [v, rot, attack] * num_robots
        rewards = np.zeros(self.num_robots)
        dones = np.zeros(self.num_robots, dtype=bool)
        infos = self._get_infos()

        # Survival rewards for alive robots
        for i, robot in enumerate(self.robots):
            if robot["alive"]:
                rewards[i] += self.survival_reward * self.dt

        # move rewards for alive robots
        for i, robot in enumerate(self.robots):
            if robot["alive"]:
                rewards[i] += (
                    self.move_reward * self.dt * np.abs(actions[3 * i])
                )  # velocity action

        # Update cooldowns and reborn
        for i, robot in enumerate(self.robots):
            if not robot["alive"]:
                robot["reborn_timer"] += self.dt
                reborn_time = robot["reborn_time"]
                if robot["reborn_timer"] >= reborn_time:
                    robot["alive"] = True
                    robot["reborn_timer"] = 0
                    robot["hp"] = robot["max_hp"]
                    robot["pos"] = self._get_spawn_pos(robot["team"], robot["type"])
                    robot["rot"] = 0
                    robot["attack_cooldown"] = 0
                    robot["bullet"] = robot["max_bullet"]
            else:
                robot["attack_cooldown"] = max(0, robot["attack_cooldown"] - self.dt)

        # Update resupply (ammo and HP) in ammo zones - gain-based rewards
        for i, robot in enumerate(self.robots):
            if robot["alive"]:
                rx, ry = int(robot["pos"][0]), int(robot["pos"][1])
                in_ammo = False
                if (
                    robot["team"] == "red"
                    and 0 <= ry < self.h
                    and 0 <= rx < self.w
                    and self.map_red_ammo[ry, rx]
                ):
                    in_ammo = True
                elif (
                    robot["team"] == "blue"
                    and 0 <= ry < self.h
                    and 0 <= rx < self.w
                    and self.map_blue_ammo[ry, rx]
                ):
                    in_ammo = True
                if in_ammo:
                    # Resupply rewards based on actual gain (scales < damage_reward)
                    bullet_reward_scale = 0.001  # Per bullet restored
                    hp_reward_scale = 0.001  # Per HP restored

                    # Resupply bullets
                    prev_bullet = robot["bullet"]
                    if robot["bullet"] < robot["max_bullet"]:
                        load_rate = robot["max_bullet"] / self.loader.bullet_load_time
                        robot["bullet"] += load_rate * self.dt
                        robot["bullet"] = min(robot["bullet"], robot["max_bullet"])
                        bullet_gain = robot["bullet"] - prev_bullet
                        rewards[i] += bullet_reward_scale * bullet_gain

                    # Heal HP
                    prev_hp = robot["hp"]
                    if robot["hp"] < robot["max_hp"]:
                        heal_rate = robot["max_hp"] / self.loader.bullet_load_time
                        robot["hp"] += heal_rate * self.dt
                        robot["hp"] = min(robot["hp"], robot["max_hp"])
                        hp_gain = robot["hp"] - prev_hp
                        rewards[i] += hp_reward_scale * hp_gain

        # Parse actions
        for i in range(self.num_robots):
            robot = self.robots[i]
            if robot["alive"]:
                v = np.clip(actions[3 * i], -1, 1) * self.max_velocity
                rot = np.clip(actions[3 * i + 1], -1, 1) * self.max_rotation
                attack = bool(actions[3 * i + 2])

                # Rotate
                robot["rot"] += rot * self.dt

                # Move (unchanged collision logic)
                dx = v * math.cos(robot["rot"]) * self.dt
                dy = v * math.sin(robot["rot"]) * self.dt
                new_pos = (
                    robot["pos"] + np.array([dx, dy]) * self.loader.resolution
                )  # since pos in px, v in m/s, res px/m

                # Improved collision detection
                valid_move = True

                # Check center
                cx, cy = int(new_pos[0]), int(new_pos[1])
                if not (
                    0 <= cx < self.w and 0 <= cy < self.h and self.map_walkable[cy, cx]
                ):
                    valid_move = False

                # Check corners
                if valid_move:
                    corners = self._get_corners(new_pos, robot["rot"])
                    for c in corners:
                        x, y = int(c[0]), int(c[1])
                        if not (
                            0 <= x < self.w
                            and 0 <= y < self.h
                            and self.map_walkable[y, x]
                        ):
                            valid_move = False
                            break

                if not valid_move:
                    # 给予撞墙惩罚
                    rewards[i] -= 0.01

                if valid_move and not self._collides_robots(new_pos, i):
                    robot["pos"] = new_pos
                    # Clamp after move just in case
                    robot["pos"][0] = np.clip(robot["pos"][0], 0, self.w - 1)
                    robot["pos"][1] = np.clip(robot["pos"][1], 0, self.h - 1)

                # Attack - added firing penalty
                if attack and robot["attack_cooldown"] <= 0 and robot["bullet"] > 0:
                    self.firing_robots.append(i)  # Mark as firing this frame
                    robot["bullet"] -= 1
                    rewards[
                        i
                    ] -= 0.001  # Bullet firing penalty (< damage reward magnitude)
                    att_r, vic_r, vic_i = self._perform_attack(i)
                    rewards[i] += att_r
                    if vic_i is not None:
                        self.damaged_robots.append(vic_i)  # Mark as damaged this frame
                        rewards[vic_i] += vic_r
                    robot["attack_cooldown"] = 1 / robot["attack_freq"]

        # Check bonus occupy - no sharing, full reward per occupier
        red_in_bonus = []
        blue_in_bonus = []
        for r in self.robots:
            if r["alive"]:
                rx, ry = int(r["pos"][0]), int(r["pos"][1])
                if 0 <= ry < self.h and 0 <= rx < self.w and self.map_bonus[ry, rx]:
                    if r["team"] == "red":
                        red_in_bonus.append(r)
                    elif r["team"] == "blue":
                        blue_in_bonus.append(r)

        red_count = len(red_in_bonus)
        blue_count = len(blue_in_bonus)
        if red_count > 0 and blue_count == 0:
            occupy_r = self.loader.reward_center_occupy * self.dt
            self.scores["red"] += occupy_r * red_count  # Scale team score with count
            for r in red_in_bonus:
                rid = self.robots.index(r)
                rewards[rid] += occupy_r  # Full reward to each (no /count)
        elif blue_count > 0 and red_count == 0:
            occupy_r = self.loader.reward_center_occupy * self.dt
            self.scores["blue"] += occupy_r * blue_count
            for r in blue_in_bonus:
                rid = self.robots.index(r)
                rewards[rid] += occupy_r

        # Time (victory rewards already scaled in init)
        self.current_time += self.dt
        done_condition = (
            self.current_time >= self.loader.time
            or max(self.scores.values()) >= self.loader.victory_reward
        )

        if done_condition:
            if self.scores["red"] != self.scores["blue"]:
                winner = "red" if self.scores["red"] > self.scores["blue"] else "blue"
                for i, r in enumerate(self.robots):
                    if r["team"] == winner:
                        rewards[i] += self.loader.victory_reward
                    else:
                        rewards[i] -= self.loader.victory_reward / 2
            dones = np.ones(self.num_robots, dtype=bool)

        # Draw scene after updates
        self._draw_scene()

        obs = self._get_obs()
        infos = self._get_infos()
        return obs, rewards, dones, False, infos  # truncated False for time

    def _get_corners(self, pos, rot):
        # Define robot corners relative to center (square robot)
        half_size = self.robot_size / 2 * self.loader.resolution
        corners_rel = np.array(
            [
                [-half_size, -half_size],
                [half_size, -half_size],
                [half_size, half_size],
                [-half_size, half_size],
            ]
        )
        # Rotate corners
        cos_r, sin_r = math.cos(rot), math.sin(rot)
        rotated = np.array(
            [[cos_r * x - sin_r * y, sin_r * x + cos_r * y] for x, y in corners_rel]
        )
        return pos + rotated

    def _collides_robots(self, new_pos, self_idx):
        half_size_px = self.robot_size / 2 * self.loader.resolution
        for j, other in enumerate(self.robots):
            if j == self_idx or not other["alive"]:
                continue
            dist = np.linalg.norm(new_pos - other["pos"])
            if dist < 2 * half_size_px:
                return True
        return False

    def _draw_scene(self):
        # Create scene image using map_vis as base
        self.scene_image = self.map_vis.copy()

        # Draw robots on scene
        for robot in self.robots:
            if robot["alive"]:
                # Draw robot as circle for simplicity in vision
                y, x = int(robot["pos"][1]), int(robot["pos"][0])
                cv2.circle(
                    self.scene_image, (x, y), self.robot_radius_px, (255, 0, 0), -1
                )

    def _render_vision(self, robot_idx):
        robot = self.robots[robot_idx]
        if not robot["alive"]:
            # Black vision for dead robots
            return np.zeros(
                (self.output_mat_size, self.output_mat_size, 3), dtype=np.uint8
            )

        # Robot position in scene
        cx, cy = int(robot["pos"][0]), int(robot["pos"][1])
        vision_radius = self.vision_size // 2

        # Crop scene around robot
        half_v = vision_radius
        scene_crop = self.scene_image[
            max(0, cy - half_v) : cy + half_v, max(0, cx - half_v) : cx + half_v
        ]

        # Resize to output size (simple bilinear)
        vision = cv2.resize(scene_crop, (self.output_mat_size, self.output_mat_size))

        # Circular mask for vision
        mask = np.zeros((self.output_mat_size, self.output_mat_size), dtype=np.uint8)
        cv2.circle(
            mask,
            (self.output_mat_size // 2, self.output_mat_size // 2),
            half_v,
            255,
            -1,
        )
        vision = cv2.bitwise_and(vision, vision, mask=mask)

        # Fill outside circle with black
        vision[mask == 0] = 0

        return vision

    def _get_obs(self):
        visions = np.array([self._render_vision(i) for i in range(self.num_robots)])
        infos = self._get_infos()
        return {"visions": visions, "infos": infos}

    def _get_infos(self):
        return {
            "time": self.current_time,
            "hp": np.array([r["hp"] for r in self.robots]),
            "bullet": np.array([r["bullet"] for r in self.robots]),
            "attack_cooldown": np.array([r["attack_cooldown"] for r in self.robots]),
            "reborn_timer": np.array([r["reborn_timer"] for r in self.robots]),
            "scores": {"red": self.scores["red"], "blue": self.scores["blue"]},
        }

    def _perform_attack(self, attacker_idx):
        robot = self.robots[attacker_idx]
        candidates = [
            r for r in self.robots if r["team"] != robot["team"] and r["alive"]
        ]

        if not candidates:
            return 0, 0, None

        # Step 1: Line of sight check using Bresenham
        visible = []
        for candidate in candidates:
            if self._line_of_sight(robot["pos"], candidate["pos"]):
                visible.append(candidate)

        if not visible:
            return 0, 0, None

        # Step 2: Select closest visible
        closest = min(visible, key=lambda c: np.linalg.norm(c["pos"] - robot["pos"]))
        target_dist_m = (
            np.linalg.norm(closest["pos"] - robot["pos"]) / self.loader.resolution
        )

        # Step 3: Compute accuracy and check hit
        accuracy = self._get_accuracy(target_dist_m)
        hit = random.random() < accuracy
        if hit:
            dmg = robot["attack_dmg"]
            damage_frac = dmg / closest["max_hp"]
            att_r = self.damage_reward * damage_frac
            vic_r = -self.damage_reward * damage_frac
            closest["hp"] -= dmg
            if closest["hp"] <= 0:
                closest["hp"] = 0
                closest["alive"] = False
                closest["reborn_timer"] = 0
                self.scores[robot["team"]] += self.loader.reward_kill
                att_r = self.loader.reward_kill
                vic_r = -self.loader.reward_kill
            victim_idx = self.robots.index(closest)
            return att_r, vic_r, victim_idx
        # Step 6: Cooldown is set in step() regardless of hit
        return 0, 0, None

    def _line_of_sight(self, p1, p2):
        # Bresenham line algorithm to check if line between p1 and p2 is walkable
        x1, y1 = int(p1[0]), int(p1[1])
        x2, y2 = int(p2[0]), int(p2[1])
        dx = abs(x2 - x1)
        dy = abs(y2 - y1)
        sx = 1 if x1 < x2 else -1
        sy = 1 if y1 < y2 else -1
        err = dx - dy

        while True:
            if not (
                0 <= y1 < self.h and 0 <= x1 < self.w and self.map_walkable[y1, x1]
            ):
                return False
            if x1 == x2 and y1 == y2:
                return True
            e2 = 2 * err
            if e2 > -dy:
                err -= dy
                x1 += sx
            if e2 < dx:
                err += dx
                y1 += sy

    def _get_accuracy(self, dist):
        keys = sorted(self.loader.attack_accuracy.keys())
        for i in range(len(keys) - 1):
            if keys[i] <= dist < keys[i + 1]:
                return self.loader.attack_accuracy[keys[i]] + (
                    self.loader.attack_accuracy[keys[i + 1]]
                    - self.loader.attack_accuracy[keys[i]]
                ) * (dist - keys[i]) / (keys[i + 1] - keys[i])
        return 0 if dist >= 20 else self.loader.attack_accuracy[keys[0]]

    def _get_spawn_pos(self, team, rtype):
        starts = {
            "red": {
                "infantry": (
                    self.loader.red_infantry_start[0]
                    if len(self.loader.red_infantry_start) > 0
                    else np.array([100.0, 100.0])
                ),
                "hero": (
                    self.loader.red_hero_start[0]
                    if len(self.loader.red_hero_start) > 0
                    else np.array([100.0, 100.0])
                ),
                "sentry": (
                    self.loader.red_sentry_start[0]
                    if len(self.loader.red_sentry_start) > 0
                    else np.array([100.0, 100.0])
                ),
            },
            "blue": {
                "infantry": (
                    self.loader.blue_infantry_start[0]
                    if len(self.loader.blue_infantry_start) > 0
                    else np.array([self.w - 100.0, self.h - 100.0])
                ),
                "hero": (
                    self.loader.blue_hero_start[0]
                    if len(self.loader.blue_hero_start) > 0
                    else np.array([self.w - 100.0, self.h - 100.0])
                ),
                "sentry": (
                    self.loader.blue_sentry_start[0]
                    if len(self.loader.blue_sentry_start) > 0
                    else np.array([self.w - 100.0, self.h - 100.0])
                ),
            },
        }
        pos = starts[team][rtype].astype(float)
        pos[0] = np.clip(pos[0], 0, self.w - 1)
        pos[1] = np.clip(pos[1], 0, self.h - 1)
        return pos

    def _adjust_brightness(self, color, factor):
        """Adjust color brightness by factor (0.7 for dim, 1.3 for bright), clamp 0-255"""
        return tuple(min(255, max(0, int(c * factor))) for c in color)

    def render(self, mode="human"):
        if self.render_mode != mode:
            self.render_mode = mode  # Update if changed

        if mode == "rgb_array":
            if self.screen is None:
                pygame.init()
                self.screen = pygame.Surface((self.w, self.h))
            if self.font is None:
                self.font = pygame.font.Font(None, 24)
        elif mode == "human":
            if self.screen is None:
                pygame.init()
                self.screen = pygame.display.set_mode((self.w, self.h))
                self.clock = pygame.time.Clock()
            if self.font is None:
                self.font = pygame.font.Font(None, 24)

        # Common drawing logic
        self.screen.fill((0, 0, 0))
        # Draw map
        surf = pygame.surfarray.make_surface(self.map_vis.swapaxes(0, 1))
        self.screen.blit(surf, (0, 0))

        # Robot params for bars
        robot_width_px = self.robot_size * self.loader.resolution
        bar_length = 1.5 * robot_width_px
        bar_height = 3
        bar_spacing = 2

        # Draw robots
        for i, robot in enumerate(self.robots):
            if not robot["alive"]:
                team_color = (128, 128, 128)
                type_color = (128, 128, 128)
            else:
                team_color_base = (255, 0, 0) if robot["team"] == "red" else (0, 0, 255)
                type_color = (
                    (0, 0, 0)
                    if robot["type"] == "infantry"
                    else (0, 255, 0) if robot["type"] == "hero" else (255, 0, 255)
                )
                # Adjust brightness based on events
                if i in self.firing_robots:
                    team_color = self._adjust_brightness(
                        team_color_base, 0.7
                    )  # Dim for firing
                elif i in self.damaged_robots:
                    team_color = self._adjust_brightness(
                        team_color_base, 1.3
                    )  # Bright for damaged
                else:
                    team_color = team_color_base

            # Draw base
            base_surf = self._get_robot_surf("base", robot, (0, 0, 0))
            rotated = pygame.transform.rotate(base_surf, -math.degrees(robot["rot"]))
            rect = rotated.get_rect(center=(robot["pos"][0], robot["pos"][1]))
            self.screen.blit(rotated, rect)
            # Armor
            armor_surf = self._get_robot_surf("armor", robot, team_color)
            rotated_a = pygame.transform.rotate(armor_surf, -math.degrees(robot["rot"]))
            rect_a = rotated_a.get_rect(center=(robot["pos"][0], robot["pos"][1]))
            self.screen.blit(rotated_a, rect_a)
            # Type mark
            mark_surf = self._get_robot_surf("type_mark", robot, type_color)
            rotated_m = pygame.transform.rotate(mark_surf, -math.degrees(robot["rot"]))
            rect_m = rotated_m.get_rect(center=(robot["pos"][0], robot["pos"][1]))
            self.screen.blit(rotated_m, rect_m)

            # Draw health and ammo bars above robot (only if alive)
            if robot["alive"]:
                center_x = robot["pos"][0]
                top_y = rect.top - bar_spacing - bar_height  # HP bar
                bottom_y = top_y - bar_spacing - bar_height  # Ammo bar above HP

                # HP bar
                hp_ratio = robot["hp"] / robot["max_hp"]
                # Filled green
                pygame.draw.rect(
                    self.screen,
                    (0, 255, 0),
                    (
                        center_x - bar_length / 2,
                        top_y,
                        bar_length * hp_ratio,
                        bar_height,
                    ),
                )
                # Empty red
                pygame.draw.rect(
                    self.screen,
                    (255, 0, 0),
                    (
                        center_x - bar_length / 2 + bar_length * hp_ratio,
                        top_y,
                        bar_length * (1 - hp_ratio),
                        bar_height,
                    ),
                )

                # Ammo bar (above HP)
                bullet_ratio = robot["bullet"] / robot["max_bullet"]
                # Filled blue
                pygame.draw.rect(
                    self.screen,
                    (0, 0, 255),
                    (
                        center_x - bar_length / 2,
                        bottom_y,
                        bar_length * bullet_ratio,
                        bar_height,
                    ),
                )
                # Empty gray
                pygame.draw.rect(
                    self.screen,
                    (128, 128, 128),
                    (
                        center_x - bar_length / 2 + bar_length * bullet_ratio,
                        bottom_y,
                        bar_length * (1 - bullet_ratio),
                        bar_height,
                    ),
                )

        # UI Elements (same for both modes)
        y_offset = 10
        line_height = 25
        # Time
        remaining_time = max(0, self.loader.time - self.current_time)
        time_text = self.font.render(f"Time: {remaining_time:.1f}s", True, (0, 255, 0))
        self.screen.blit(time_text, (10, y_offset))
        y_offset += line_height
        # Scores
        score_text = self.font.render(
            f"Red: {self.scores['red']} | Blue: {self.scores['blue']}",
            True,
            (0, 255, 0),
        )
        self.screen.blit(score_text, (10, y_offset))
        y_offset += line_height
        # Robot status header
        robot_info_text = self.font.render("Robot Status:", True, (0, 255, 0))
        self.screen.blit(robot_info_text, (10, y_offset))
        y_offset += line_height
        # Per robot
        for robot in self.robots:
            team_color = (255, 0, 0) if robot["team"] == "red" else (0, 0, 255)
            type_abbr = robot["type"][0].upper()
            if robot["alive"]:
                hp_text = f"{robot['team'].capitalize()} {type_abbr}: HP {robot['hp']}/{robot['max_hp']} Bullet {robot['bullet']}/{robot['max_bullet']}"
                text_color = team_color
            else:
                remaining_reborn = max(0, robot["reborn_time"] - robot["reborn_timer"])
                hp_text = f"{robot['team'].capitalize()} {type_abbr}: Dead (Reborn: {remaining_reborn:.1f}s)"
                text_color = (128, 128, 128)
            robot_text = self.font.render(hp_text, True, text_color)
            self.screen.blit(robot_text, (10, y_offset))
            y_offset += line_height

        if mode == "human":
            pygame.display.flip()
            self.clock.tick(self.metadata["render_fps"])
            return None
        elif mode == "rgb_array":
            # Convert to (H, W, 3) RGB array
            arr = pygame.surfarray.array3d(self.screen)
            arr = np.transpose(arr, axes=(1, 0, 2))  # (W, H, 3) -> (H, W, 3)
            return arr

    def _get_robot_surf(self, part, robot, color=(0, 255, 0)):
        if part == "base":
            mask = self.loader.robot_base["general"]
        elif part == "armor":
            mask = self.loader.robot_armor["general"]
        else:  # type_mark
            mask = self.loader.robot_type_mark["general"]

        mask_sq = mask.squeeze()  # (h, w)
        h, w = mask_sq.shape
        alpha = mask_sq.astype(np.float32) / 255.0
        rgb = np.full((h, w, 3), list(color), dtype=np.uint8)
        a = (alpha * 255).astype(np.uint8)

        # Create surface with per-pixel alpha
        surf = pygame.Surface((w, h), pygame.SRCALPHA)
        pixels3d = pygame.surfarray.pixels3d(surf)
        pixels3d[:] = rgb.transpose(1, 0, 2)  # (h,w,3) -> (w,h,3)
        del pixels3d

        pixels_alpha = pygame.surfarray.pixels_alpha(surf)
        pixels_alpha[:] = a.transpose()  # (h,w) -> (w,h)
        del pixels_alpha

        return surf

    def close(self):
        if self.screen is not None:
            pygame.quit()


# Main for testing with random actions and visualization
if __name__ == "__main__":
    env = rm_simple_simulator(render_mode="human")
    obs, _ = env.reset()
    done = False
    random_idx = random.randint(0, env.num_robots - 1)
    while not done:
        actions = env.action_space.sample()  # random
        obs, rewards, terminated, truncated, info = env.step(actions)
        done = terminated.any() or truncated
        env.render()

        # Visualize random robot's vision with cv2
        vis = obs["visions"][random_idx]
        vis_bgr = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)
        cv2.imshow(f"Vision of robot {random_idx}", vis_bgr)
        key = cv2.waitKey(1) & 0xFF
        if key == ord("q"):
            done = True

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                done = True
    env.close()
    cv2.destroyAllWindows()
